import re
import sys
import json
from typing import Dict, TypedDict
from langgraph.graph import END, StateGraph
from doubao import doubao_chat

# 大模型的密钥
API_KEY = sys.argv[1]


# 定义状态结构
class AgentState(TypedDict):
    user_input: str  # 用户输入的字符串
    a: int           # 解析出的第一个数字
    b: int           # 解析出的第二个数字
    operation: str   # 解析出的操作类型
    result: int      # 计算结果


def llm_chat(state: AgentState) -> Dict:
    """使用豆包大模型解析用户输入，提取数字和操作类型"""
    print(f"\n🧠 用户输入: '{state['user_input']}'")

    # 构建给大模型的提示
    prompt = f"""
    请从以下用户问题中提取两个整数和要执行的操作类型：
    "{state['user_input']}"

    输出格式必须是纯JSON格式，包含且仅包含以下三个字段：
    - "a": 第一个整数
    - "b": 第二个整数
    - "operation": 操作类型（只能是 "add" 或 "subtract"）

    示例输出：{{"a": 5, "b": 3, "operation": "add"}}

    只输出JSON，不要包含任何其他内容。
    """

    # 调用豆包大模型API
    llm_output = doubao_chat(api_key=API_KEY, text=prompt)

    # 提取模型返回的内容
    content = llm_output.message.content
    print(f"🤖 豆包大模型原始输出: {content}")
    if not content:
        raise ValueError("模型返回内容为空")

    # 从输出中提取JSON部分（模型可能会添加额外说明）
    json_match = re.search(r'\{.*\}', content, re.DOTALL)
    if not json_match:
        raise ValueError("未找到有效的JSON输出")

    # 将JSON字符串转换为字典
    result_dict = json.loads(json_match.group(0))

    # 验证结果
    if 'a' not in result_dict or 'b' not in result_dict or 'operation' not in result_dict:
        raise ValueError("返回的JSON缺少必要字段")

    # 转换数字类型
    a = int(result_dict['a'])
    b = int(result_dict['b'])
    operation = result_dict['operation']

    print(f"🔍 豆包大模型解析: 数字 {a} 和 {b}, 操作: {operation}")
    return {"a": a, "b": b, "operation": operation}


# 计算节点函数
def add_numbers(state: AgentState) -> Dict[str, int]:
    result = state["a"] + state["b"]
    print(f"➕ 计算结果: {state['a']} + {state['b']} = {result}")
    return {"result": result}


def subtract_numbers(state: AgentState) -> Dict[str, int]:
    result = state["a"] - state["b"]
    print(f"➖ 计算结果: {state['a']} - {state['b']} = {result}")
    return {"result": result}


# 结果展示节点
def display_result(state: AgentState) -> Dict:
    print(f"✅ 最终结果: {state['result']}")
    return {}


# 构建图工作流
def create_workflow():
    # 初始化状态图
    workflow = StateGraph(AgentState)

    # 添加节点
    workflow.add_node("chat", llm_chat)
    workflow.add_node("add", add_numbers)
    workflow.add_node("subtract", subtract_numbers)
    workflow.add_node("display", display_result)

    # 设置入口点
    workflow.set_entry_point("chat")

    # 根据解析结果路由
    workflow.add_conditional_edges(
        "chat",
        lambda state: state["operation"],
        {
            "add": "add",
            "subtract": "subtract",
        }
    )

    # 连接计算节点到显示节点
    workflow.add_edge("add", "display")
    workflow.add_edge("subtract", "display")

    # 连接到结束
    workflow.add_edge("display", END)

    return workflow.compile()


# 执行示例
if __name__ == "__main__":
    # 初始化工作流
    app = create_workflow()

    # 测试案例
    test_cases = [
        "5加3等于多少？",
        "10减去7",
        "计算20和15的和",
        "100 - 50",
        "8和4的差是多少？",
        "简单的计算：3和5",  # 测试默认操作
        "只有数字：42 100",  # 测试无操作符
    ]

    for input_str in test_cases:
        print("\n" + "="*50)
        print(f"测试输入: '{input_str}'")
        result = app.invoke({"user_input": input_str})
        print(f"完整输出: {result}")
        break
